import os
import numpy as np
import pandas as pd
import logging
import xarray as xr
from scipy.interpolate import griddata
from global_land_mask import globe
from tqdm import tqdm

# 配置参数
INPUT_CSV = "/mnt/datastore/liudddata/cloudsat_data/merged_201902_202004_cloudsat.csv"
# OUTPUT_NC = "/mnt/datastore/liudddata/cloudsat_data/cloudsat_seasonal35_mean_2019MAM.nc"
OUTPUT_NC = "/mnt/datastore/liudddata/cloudsat_data/cloudsat_annual_mean_2019MAM.nc"
FY4A_GRID = "/home/liudd/data_preprocessing/FY4A_coordinates.nc"

# 初始化日志
logging.basicConfig(
    filename='cloudsat_processing.log',
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)


def load_fy4a_grid():
    """加载FY4A网格并生成海洋掩膜"""
    try:
        ds = xr.open_dataset(FY4A_GRID)
        lon_2d = ds.lon.values.T
        lat_2d = ds.lat.values.T

        # 检查纬度范围
        lat_2d = np.clip(lat_2d, -90, 90)

        # 生成海洋掩膜
        valid_mask = np.isfinite(lon_2d) & np.isfinite(lat_2d)
        lon_valid = np.where(lon_2d[valid_mask] > 180,
                             lon_2d[valid_mask] - 360,
                             lon_2d[valid_mask])
        is_land = globe.is_land(lat_2d[valid_mask], lon_valid)

        ocean_mask = np.full(lon_2d.shape, False)
        ocean_mask[valid_mask] = ~is_land

        return {
            'lon': lon_2d,
            'lat': lat_2d,
            'ocean_mask': ocean_mask,
            'shape': lon_2d.shape
        }
    except Exception as e:
        logging.error(f"加载FY4A网格时出错: {str(e)}")
        raise


def seasonal_average(df):
    """计算季节平均"""
    try:
        # 筛选2019年3-5月数据
        df['time'] = pd.to_datetime(df['time'])
        mask = (df['time'].dt.year == 2019) & (df['time'].dt.month.between(2, 12))
        # mask = (df['time'].dt.year == 2019) & (df['time'].dt.month.between(2, 12)) | (df['time'].dt.year == 2020) & (
        #             df['time'].dt.month == 1)
        seasonal_df = df[mask].copy()

        # 处理缺失值
        seasonal_df = seasonal_df.dropna(subset=['cloudsat_cbh'])

        # # 空间分箱平均（1°×1°网格）
        # seasonal_df['lon_bin'] = np.floor(seasonal_df['longitude'])
        # seasonal_df['lat_bin'] = np.floor(seasonal_df['latitude'])
        # 空间分箱平均（0.5°×0.5°网格）
        seasonal_df['lon_bin'] = np.floor(seasonal_df['longitude'] / 0.02) * 0.02
        seasonal_df['lat_bin'] = np.floor(seasonal_df['latitude'] / 0.02) * 0.02

        # 分组计算
        grouped = seasonal_df.groupby(['lon_bin', 'lat_bin'])
        avg_df = grouped['cloudsat_cbh'].mean().reset_index()

        return avg_df[['lon_bin', 'lat_bin', 'cloudsat_cbh']]
    except Exception as e:
        logging.error(f"计算季节平均时出错: {str(e)}")
        raise


def interpolate_to_grid(avg_df, grid_info):
    """将季节平均插值到FY4A网格"""
    try:
        # 准备插值数据
        points = avg_df[['lon_bin', 'lat_bin']].values
        values = avg_df['cloudsat_cbh'].values

        # 生成目标网格
        target_lon = grid_info['lon'].ravel()
        target_lat = grid_info['lat'].ravel()
        valid = ~np.isnan(target_lon) & ~np.isnan(target_lat)

        # 分块插值（内存优化）
        chunk_size = 1000000
        interpolated = np.full(target_lon.shape, np.nan)

        for i in tqdm(range(0, len(target_lon), chunk_size),
                      desc="插值进度"):
            chunk_slice = slice(i, min(i + chunk_size, len(target_lon)))
            chunk_valid = valid[chunk_slice]

            if np.any(chunk_valid):
                chunk_points = np.column_stack(
                    (target_lon[chunk_slice][chunk_valid],
                     target_lat[chunk_slice][chunk_valid])
                )

                chunk_values = griddata(
                    points, values,
                    chunk_points,
                    # method='linear',
                    method='nearest',  # 改为邻近插值
                    fill_value=np.nan
                )

                interpolated[chunk_slice][chunk_valid] = chunk_values

        return interpolated.reshape(grid_info['shape'])
    except Exception as e:
        logging.error(f"网格插值时出错: {str(e)}")
        raise


def apply_ocean_mask(data, grid_info):
    """应用海洋掩膜"""
    try:
        return np.where(grid_info['ocean_mask'], data, np.nan)
    except Exception as e:
        logging.error(f"应用海洋掩膜时出错: {str(e)}")
        raise


def save_netcdf(data, grid_info):
    """保存结果到NetCDF"""
    try:
        ds = xr.Dataset(
            {
                "cbh": (("y", "x"), data),
                "ocean_mask": (("y", "x"), grid_info['ocean_mask'])
            },
            coords={
                "lon": (("y", "x"), grid_info['lon']),
                "lat": (("y", "x"), grid_info['lat'])
            },
            attrs={
                "description": "2019年春季(3-5月)平均云底高度",
                "units": "米",
                "processing_method": "空间分箱平均后线性插值",
                "time_range": "2019年2月 - 2020年1月",
                "data_source": INPUT_CSV
            }
        )
        ds.to_netcdf(OUTPUT_NC)
    except Exception as e:
        logging.error(f"保存NetCDF文件时出错: {str(e)}")
        raise


def main():
    logging.info("=== 开始处理流程 ===")

    try:
        # 步骤1: 加载并处理数据
        df = pd.read_csv(INPUT_CSV)
        logging.info(f"原始数据量: {len(df)} 条")

        # 步骤2: 计算季节平均
        avg_df = seasonal_average(df)
        logging.info(f"分箱平均点数: {len(avg_df)} 个")

        # 步骤3: 加载FY4A网格
        grid_info = load_fy4a_grid()

        # 步骤4: 网格插值
        interpolated_data = interpolate_to_grid(avg_df, grid_info)

        # 步骤5: 应用海洋掩膜
        final_data = apply_ocean_mask(interpolated_data, grid_info)

        # 步骤6: 保存结果
        save_netcdf(final_data, grid_info)
        logging.info("处理流程完成")
    except Exception as e:
        logging.critical(f"主流程失败: {str(e)}")
        raise


if __name__ == "__main__":
    try:
        main()
        print("处理成功完成!")
    except Exception as e:
        print(f"处理失败: {str(e)}")